// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "net/base/keygen_handler.h"

#include <openssl/bytestring.h>
#include <openssl/evp.h>
#include <stdint.h>

#include <string>
#include <utility>

#include "base/base64.h"
#include "base/bind.h"
#include "base/location.h"
#include "base/logging.h"
#include "base/strings/string_piece.h"
#include "base/synchronization/waitable_event.h"
#include "base/threading/thread_restrictions.h"
#include "base/threading/worker_pool.h"
#include "build/build_config.h"
#include "crypto/scoped_openssl_types.h"
#include "testing/gtest/include/gtest/gtest.h"

#if defined(USE_NSS_CERTS)
#include "crypto/nss_crypto_module_delegate.h"
#include "crypto/scoped_test_nss_db.h"
#include <private/pprthred.h> // PR_DetachThread
#endif

namespace net {

namespace {

#if defined(USE_NSS_CERTS)
    class StubCryptoModuleDelegate : public crypto::NSSCryptoModuleDelegate {
    public:
        explicit StubCryptoModuleDelegate(crypto::ScopedPK11Slot slot)
            : slot_(std::move(slot))
        {
        }

        std::string RequestPassword(const std::string& slot_name,
            bool retry,
            bool* cancelled) override
        {
            return std::string();
        }

        crypto::ScopedPK11Slot RequestSlot() override
        {
            return crypto::ScopedPK11Slot(PK11_ReferenceSlot(slot_.get()));
        }

    private:
        crypto::ScopedPK11Slot slot_;
    };
#endif

    const char kChallenge[] = "some challenge";

    class KeygenHandlerTest : public ::testing::Test {
    public:
        KeygenHandlerTest() { }
        ~KeygenHandlerTest() override { }

        std::unique_ptr<KeygenHandler> CreateKeygenHandler()
        {
            std::unique_ptr<KeygenHandler> handler(
                new KeygenHandler(768, kChallenge, GURL("http://www.example.com")));
#if defined(USE_NSS_CERTS)
            handler->set_crypto_module_delegate(
                std::unique_ptr<crypto::NSSCryptoModuleDelegate>(
                    new StubCryptoModuleDelegate(crypto::ScopedPK11Slot(
                        PK11_ReferenceSlot(test_nss_db_.slot())))));
#endif
            return handler;
        }

    private:
#if defined(USE_NSS_CERTS)
        crypto::ScopedTestNSSDB test_nss_db_;
#endif
    };

    base::StringPiece StringPieceFromCBS(const CBS& cbs)
    {
        return base::StringPiece(reinterpret_cast<const char*>(CBS_data(&cbs)),
            CBS_len(&cbs));
    }

    // Assert that |result| is a valid output for KeygenHandler given challenge
    // string of |challenge|.
    void AssertValidSignedPublicKeyAndChallenge(const std::string& result,
        const std::string& challenge)
    {
        // Verify it's valid base64:
        std::string spkac;
        ASSERT_TRUE(base::Base64Decode(result, &spkac));

        // Parse the following structure:
        //
        //   PublicKeyAndChallenge ::= SEQUENCE {
        //     spki SubjectPublicKeyInfo,
        //     challenge IA5STRING
        //   }
        //   SignedPublicKeyAndChallenge ::= SEQUENCE {
        //     publicKeyAndChallenge PublicKeyAndChallenge,
        //     signatureAlgorithm AlgorithmIdentifier,
        //     signature BIT STRING
        //   }

        CBS cbs;
        CBS_init(&cbs, reinterpret_cast<const uint8_t*>(spkac.data()), spkac.size());

        // The input should consist of a SEQUENCE.
        CBS child;
        ASSERT_TRUE(CBS_get_asn1(&cbs, &child, CBS_ASN1_SEQUENCE));
        ASSERT_EQ(0u, CBS_len(&cbs));

        // Extract the raw PublicKeyAndChallenge.
        CBS public_key_and_challenge_raw;
        ASSERT_TRUE(CBS_get_asn1_element(&child, &public_key_and_challenge_raw,
            CBS_ASN1_SEQUENCE));

        // Parse out the PublicKeyAndChallenge.
        CBS copy = public_key_and_challenge_raw;
        CBS public_key_and_challenge;
        ASSERT_TRUE(
            CBS_get_asn1(&copy, &public_key_and_challenge, CBS_ASN1_SEQUENCE));
        ASSERT_EQ(0u, CBS_len(&copy));
        crypto::ScopedEVP_PKEY key(EVP_parse_public_key(&public_key_and_challenge));
        ASSERT_TRUE(key);
        CBS challenge_spkac;
        ASSERT_TRUE(CBS_get_asn1(&public_key_and_challenge, &challenge_spkac,
            CBS_ASN1_IA5STRING));
        ASSERT_EQ(0u, CBS_len(&public_key_and_challenge));

        // The challenge must match.
        ASSERT_EQ(challenge, StringPieceFromCBS(challenge_spkac));

        // The next element must be the AlgorithmIdentifier for MD5 with RSA.
        static const uint8_t kMd5WithRsaEncryption[] = {
            0x30,
            0x0d,
            0x06,
            0x09,
            0x2a,
            0x86,
            0x48,
            0x86,
            0xf7,
            0x0d,
            0x01,
            0x01,
            0x04,
            0x05,
            0x00,
        };
        CBS algorithm;
        ASSERT_TRUE(CBS_get_bytes(&child, &algorithm, sizeof(kMd5WithRsaEncryption)));
        ASSERT_EQ(
            base::StringPiece(reinterpret_cast<const char*>(kMd5WithRsaEncryption),
                sizeof(kMd5WithRsaEncryption)),
            StringPieceFromCBS(algorithm));

        // Finally, parse the signature.
        CBS signature;
        ASSERT_TRUE(CBS_get_asn1(&child, &signature, CBS_ASN1_BITSTRING));
        ASSERT_EQ(0u, CBS_len(&child));
        uint8_t pad;
        ASSERT_TRUE(CBS_get_u8(&signature, &pad));
        ASSERT_EQ(0u, pad);

        // Check the signature.
        crypto::ScopedEVP_MD_CTX ctx(EVP_MD_CTX_create());
        ASSERT_TRUE(
            EVP_DigestVerifyInit(ctx.get(), nullptr, EVP_md5(), nullptr, key.get()));
        ASSERT_TRUE(EVP_DigestVerifyUpdate(ctx.get(),
            CBS_data(&public_key_and_challenge_raw),
            CBS_len(&public_key_and_challenge_raw)));
        ASSERT_TRUE(EVP_DigestVerifyFinal(ctx.get(), CBS_data(&signature),
            CBS_len(&signature)));
    }

    TEST_F(KeygenHandlerTest, SmokeTest)
    {
        std::unique_ptr<KeygenHandler> handler(CreateKeygenHandler());
        handler->set_stores_key(false); // Don't leave the key-pair behind
        std::string result = handler->GenKeyAndSignChallenge();
        VLOG(1) << "KeygenHandler produced: " << result;
        AssertValidSignedPublicKeyAndChallenge(result, kChallenge);
    }

    void ConcurrencyTestCallback(const std::string& challenge,
        base::WaitableEvent* event,
        std::unique_ptr<KeygenHandler> handler,
        std::string* result)
    {
        // We allow Singleton use on the worker thread here since we use a
        // WaitableEvent to synchronize, so it's safe.
        base::ThreadRestrictions::ScopedAllowSingleton scoped_allow_singleton;
        handler->set_stores_key(false); // Don't leave the key-pair behind.
        *result = handler->GenKeyAndSignChallenge();
        event->Signal();
#if defined(USE_NSS_CERTS)
        // Detach the thread from NSPR.
        // Calling NSS functions attaches the thread to NSPR, which stores
        // the NSPR thread ID in thread-specific data.
        // The threads in our thread pool terminate after we have called
        // PR_Cleanup.  Unless we detach them from NSPR, net_unittests gets
        // segfaults on shutdown when the threads' thread-specific data
        // destructors run.
        PR_DetachThread();
#endif
    }

    // We asynchronously generate the keys so as not to hang up the IO thread. This
    // test tries to catch concurrency problems in the keygen implementation.
    TEST_F(KeygenHandlerTest, ConcurrencyTest)
    {
        const int NUM_HANDLERS = 5;
        base::WaitableEvent* events[NUM_HANDLERS] = { NULL };
        std::string results[NUM_HANDLERS];
        for (int i = 0; i < NUM_HANDLERS; i++) {
            std::unique_ptr<KeygenHandler> handler(CreateKeygenHandler());
            events[i] = new base::WaitableEvent(
                base::WaitableEvent::ResetPolicy::AUTOMATIC,
                base::WaitableEvent::InitialState::NOT_SIGNALED);
            base::WorkerPool::PostTask(FROM_HERE,
                base::Bind(ConcurrencyTestCallback,
                    "some challenge",
                    events[i],
                    base::Passed(&handler),
                    &results[i]),
                true);
        }

        for (int i = 0; i < NUM_HANDLERS; i++) {
            // Make sure the job completed
            events[i]->Wait();
            delete events[i];
            events[i] = NULL;

            VLOG(1) << "KeygenHandler " << i << " produced: " << results[i];
            AssertValidSignedPublicKeyAndChallenge(results[i], "some challenge");
        }
    }

} // namespace

} // namespace net
